今天是第五天，前面幾天介紹lstm範例，今天可以寫一個簡單的yolo來分析多隻斑馬魚的行為，以下是程式碼

import cv2
import numpy as np

# 載入 YOLO 模型
def load_yolo_model():
    net = cv2.dnn.readNet("yolov3.weights", "yolov3.cfg")
    with open("coco.names", "r") as f:
        classes = [line.strip() for line in f.readlines()]
    layer_names = net.getLayerNames()
    output_layers = [layer_names[i - 1] for i in net.getUnconnectedOutLayers()]
    return net, classes, output_layers

# 偵測斑馬魚
def detect_fish(frame, net, output_layers):
    height, width, channels = frame.shape
    blob = cv2.dnn.blobFromImage(frame, 0.00392, (416, 416), (0, 0, 0), True, crop=False)
    net.setInput(blob)
    outs = net.forward(output_layers)

    class_ids = []
    confidences = []
    boxes = []

    for out in outs:
        for detection in out:
            scores = detection[5:]
            class_id = np.argmax(scores)
            confidence = scores[class_id]
            if confidence > 0.5 and class_id == 0:  # 假設斑馬魚的類別ID是0
                center_x = int(detection[0] * width)
                center_y = int(detection[1] * height)
                w = int(detection[2] * width)
                h = int(detection[3] * height)
                x = int(center_x - w / 2)
                y = int(center_y - h / 2)
                boxes.append([x, y, w, h])
                confidences.append(float(confidence))
                class_ids.append(class_id)

    indexes = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.4)
    return boxes, confidences, indexes

# 繪製偵測結果
def draw_labels(boxes, confidences, indexes, frame):
    font = cv2.FONT_HERSHEY_PLAIN
    for i in range(len(boxes)):
        if i in indexes:
            x, y, w, h = boxes[i]
            label = str(confidences[i])
            color = (0, 255, 0)
            cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)
            cv2.putText(frame, label, (x, y - 10), font, 1, color, 1)
    return frame

# 主程式
def main():
    net, classes, output_layers = load_yolo_model()
    cap = cv2.VideoCapture("zebrafish_video.mp4")

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        boxes, confidences, indexes = detect_fish(frame, net, output_layers)
        frame = draw_labels(boxes, confidences, indexes, frame)

        cv2.imshow("Zebrafish Detection", frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()

1. 載入 YOLO 模型 (load_yolo_model 函數)：

   • 這個函數載入 YOLO 的權重 (yolov3.weights) 和配置文件 (yolov3.cfg)。
   • 它還讀取 coco.names 文件，將每個類別名稱儲存到 classes 列表中。
   • 獲取 YOLO 模型的輸出層名稱，並返回網絡模型、類別名稱和輸出層。

2. 偵測斑馬魚 (detect_fish 函數)：

   • 函數接收一個影像框架、YOLO 網絡模型及其輸出層。
   • 將影像轉換為 YOLO 需要的輸入格式（blob），並進行前向傳播以獲取偵測結果。
   • 解析 YOLO 的輸出，篩選出信心度超過 0.5 且類別 ID 為 0 的偵測結果（假設斑馬魚的類別 ID 是 0）。
   • 儲存偵測結果的邊界框、信心度及類別 ID。

3. 繪製偵測結果 (draw_labels 函數)：

   • 函數接收偵測到的邊界框、信心度、索引及影像框架。
   • 使用 OpenCV 在每個偵測到的斑馬魚周圍繪製邊界框，並標示信心度。

4. 主函數 (main 函數)：

   • 載入 YOLO 模型和類別名稱。
   • 打開視頻檔案 (zebrafish_video.mp4) 並逐幀讀取。
   • 對每個影像框架進行斑馬魚偵測，並繪製偵測結果。
   • 顯示處理後的影像，按下 'q' 鍵結束程式。

整個流程是：從影片中逐幀讀取影像，將影像輸入 YOLO 模型進行偵測，並將結果可視化顯示在影像上。這樣可以實時追蹤影片中的多隻斑馬魚。